"""Data collection module for quantum cloud platforms.

This module provides a unified interface for submitting quantum circuits and
collecting results from different quantum computing platforms (Tianyan and Quafu).
"""

import os
import re
import logging
from typing import Union, List, Optional
import json

class GetData:
    """Handler for quantum cloud platform data collection.
    
    This class provides a unified interface for submitting circuits and collecting
    results from different quantum computing platforms.
    
    Attributes:
        platform_name: Name of quantum platform ("tianyan" or "quafu")
        login_key: Authentication key for platform
        machine_name: Name of quantum device to use
        shots: Number of shots per circuit
        num_qubits: Number of qubits in circuits
        exp_name: Experiment name
        platform: Platform-specific instance
        _data: Platform-specific data handler
    """

    def __init__(self, 
                 platform: str,
                 login_key: str,
                 machine_name: str,
                 shots: int,
                 num_qubits: int,
                 exp_name: str):
        """Initialize platform connection and data handler.

        Args:
            platform: Platform name ("tianyan" or "quafu")
            login_key: Authentication key for platform
            machine_name: Name of quantum device
            shots: Number of shots per circuit
            num_qubits: Number of qubits in circuits
            exp_name: Experiment name
        """
        self.platform_name = platform
        self.login_key = login_key
        self.machine_name = machine_name
        self.shots = shots
        self.num_qubits = num_qubits
        self.exp_name = exp_name

        # Initialize platform-specific components
        if self.platform_name == "tianyan":
            from cqlib import TianYanPlatform
            from .cloud.tianyan_data import TianyanData
            
            self.platform = TianYanPlatform(login_key=self.login_key)
            self.platform.set_machine(self.machine_name)
            self._data = TianyanData(
                self.num_qubits, 
                self.shots, 
                self.platform, 
                self.exp_name
            )
        
        elif self.platform_name == "quafu":
            from quafu import Task, User
            from .cloud.quafu_data import QuafuData
            
            user = User(self.login_key)
            user.save_apitoken()
            self.platform = Task()
            self.platform.config(
                backend=self.machine_name, 
                shots=self.shots, 
                compile=False
            )
            self._data = QuafuData(
                self.num_qubits, 
                self.shots, 
                self.platform, 
                self.exp_name
            )

    def check_valid_gates(self) -> None:
        """Check and print valid gates for the current platform."""
        if self.platform_name == "tianyan":
            pass  # TODO: Implement for Tianyan
        elif self.platform_name == "quafu":
            gates = self.platform.backend.get_valid_gates()
            gates = [re.findall(r'"(.*?)"', gate_str)[0] for gate_str in gates]
            print(f"Valid gates for {self.platform_name}: {gates}")

    def submit_experiment(self, qasm_circuit: Union[str, List[str]], benchmark: bool = False) -> None:
        """Submit quantum circuit experiment to the cloud platform.

        Args:
            qasm_circuit: QASM format quantum circuit(s). Can be either:
                - str: Single QASM circuit string
                - List[str]: List of QASM circuit strings for batch execution
            benchmark: Whether to use benchmark mode
        """
        def num2bit_string(num_bits: int, num: int) -> str:
            """Convert a number to a binary string with specified length.
            
            Args:
                num_bits: Length of the binary string
                num: Number to convert
                
            Returns:
                Binary string representation padded to num_bits length
            """
            return format(num, f'0{num_bits}b')

        # Convert single circuit to list format for unified processing
        if isinstance(qasm_circuit, str):
            qasm_circuit = [qasm_circuit]

        if self.platform_name == "tianyan":
            from cqlib.utils import QasmToQcis
            
            from cqlib.circuits import Circuit

            circuits = []
            for i, qasm in enumerate(qasm_circuit):
                self._data.df.loc[i, 'qasm2'] = qasm
                # Extract measured qubits from qasm by finding measure instructions
                measured_qubits = [int(q.split('[')[1].split(']')[0]) 
                                 for q in qasm.split('\n') 
                                 if q.strip().startswith('measure')]
                measured_cbits = [int(q.split('->')[1].split('[')[1].split(']')[0])
                                 for q in qasm.split('\n')
                                 if q.strip().startswith('measure')]
                self._data.df.loc[i, 'qubits'] = json.dumps(measured_qubits)
                self._data.df.loc[i, 'cbits'] = json.dumps(measured_cbits)

                if benchmark:
                    self._data.df.loc[i, 'true_value'] = num2bit_string(self.num_qubits, i)

                # Convert QASM to QCIS and fix PI precision if needed
                qcis_str = QasmToQcis().convert_to_qcis(qasm)
                if "3.141593" in qcis_str:
                    qcis_str = qcis_str.replace("3.141593", "3.1415926")
                
                circ = Circuit.load(qcis_str)
                circuits.append(circ)

            self._data.submit_experiment(circuits)

        elif self.platform_name == "quafu":
            from quafu import QuantumCircuit
            
            circuits = []
            for i, qasm in enumerate(qasm_circuit):
                self._data.df.loc[i, 'qasm2'] = qasm
                # Extract measured qubits from qasm by finding measure instructions
                measured_qubits = [int(q.split('[')[1].split(']')[0]) 
                                 for q in qasm.split('\n') 
                                 if q.strip().startswith('measure')]
                measured_cbits = [int(q.split('->')[1].split('[')[1].split(']')[0])
                                 for q in qasm.split('\n')
                                 if q.strip().startswith('measure')]
                self._data.df.loc[i, 'qubits'] = json.dumps(measured_qubits)
                self._data.df.loc[i, 'cbits'] = json.dumps(measured_cbits)

                if benchmark:
                    self._data.df.loc[i, 'true_value'] = num2bit_string(self.num_qubits, i)

                circ = QuantumCircuit(self.num_qubits)
                circ.from_openqasm(qasm)
                circuits.append(circ)

            self._data.submit_experiment(circuits)
        
        self.experiment_submitted = True
        self._save_temp_backup()
    def receive_results(self) -> None:
        """Try to receive experiment results without waiting."""
        if not hasattr(self, 'experiment_submitted'):
            logging.error("No experiment has been submitted yet")
            return

        logging.basicConfig(level=logging.INFO)
        logging.info("Attempting to retrieve experiment results")

        try:
            self._data.receive_result()
            self._cleanup_temp_file()
            logging.info("Successfully retrieved experiment results")
        except Exception as e:
            logging.warning(f"Results not ready yet: {str(e)}. Please try again later.")
            # Keep temp file for next attempt
            logging.info("Temporary backup file preserved for next attempt")
    
    def save_results(self, file_name: str, simplify: bool = True) -> None:
        """Save experiment results to CSV file.

        Args:
            file_name: Path to output CSV file
            simplify: If True, save only essential columns
        """
        if simplify:
            columns_to_save = ['qubits', 'cbits', 'true_value', 'exp_result', 'qasm2']
            self._data.df[columns_to_save].to_csv(file_name, index=False)
        else:
            self._data.df.to_csv(file_name, index=False)

    def _save_temp_backup(self) -> None:
        """Save temporary backup of experiment data."""
        self.temp_file = "temp.csv"
        self.save_results(self.temp_file, simplify=False)
        logging.info(f"Saved temporary backup to {self.temp_file}")

    def _cleanup_temp_file(self) -> None:
        """Clean up temporary backup file."""
        if hasattr(self, 'temp_file') and os.path.exists(self.temp_file):
            try:
                os.remove(self.temp_file)
                logging.info(f"Deleted temporary file: {self.temp_file}")
            except OSError as e:
                logging.warning(f"Failed to delete temporary file: {str(e)}")


class ShotTestReceiver:
    """Handler for testing quantum circuits with different shot counts.
    
    This class manages submitting and collecting results for quantum circuits
    executed with varying numbers of shots.
    
    Attributes:
        platform_config: Configuration dict for quantum platform
        shot_range: List of shot counts to test
        circuit_nums: Number of circuits to run
        data_receivers: List of GetData instances for each shot count
        benchmark_flag: Whether to run benchmark circuits
    """

    def __init__(self, platform_config: dict, shot_range: list, circuit_nums: int, 
                 benchmark_flag: bool = False):
        """Initialize shot test handler.

        Args:
            platform_config: Platform configuration dictionary
            shot_range: List of shot counts to test
            circuit_nums: Number of circuits to run
            benchmark_flag: Whether to run benchmark circuits
        """
        self.platform_config = platform_config
        self.shot_range = shot_range
        self.circuit_nums = circuit_nums
        self.data_receivers = []
        self.benchmark_flag = benchmark_flag
        
        # Create GetData instance for each shot count
        for shots in shot_range:
            self.data_receivers.append(
                GetData(platform_config['platform'],
                       platform_config['login_key'],
                       platform_config['machine_name'],
                       shots,
                       platform_config['num_qubits'],
                       platform_config['exp_name']))

    def submit_experiment(self, target_circuit: Union[str, List[str]]) -> None:
        """Submit quantum circuits for execution.

        Args:
            target_circuit: Circuit(s) to execute, either single circuit or list
        """
        if isinstance(target_circuit, str):
            target_circuit = [target_circuit] * self.circuit_nums
            
        for data_receiver in self.data_receivers:
            data_receiver.submit_experiment(target_circuit, 
                                         benchmark=self.benchmark_flag)

    def receive_results(self) -> None:
        """Collect results from all experiments."""
        for data_receiver in self.data_receivers:
            data_receiver.receive_results()

    def save_results(self, save_path: str, simplify: bool = True) -> None:
        """Save results from all experiments.

        Args:
            save_path: Directory to save results
            simplify: Whether to save only essential columns
        """
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        for data_receiver in self.data_receivers:
            # Generate base filename
            base_name = f'{self.platform_config["exp_name"]}_{data_receiver.shots}shots'
            file_name = f'{base_name}.csv'
            file_path = os.path.join(save_path, file_name)
            
            # Check if file exists and append number if needed
            counter = 1
            while os.path.exists(file_path):
                file_name = f'{base_name}_{counter}.csv'
                file_path = os.path.join(save_path, file_name)
                counter += 1
                
            data_receiver.save_results(file_path, simplify)


class CircuitReceiver:
    """Handler for executing quantum circuits with fixed shot count.
    
    This class manages submitting and collecting results for quantum circuits
    executed with a single shot count.
    
    Attributes:
        platform_config: Configuration dict for quantum platform
        circuit_nums: Number of circuits to run
        shots: Number of shots per circuit
        benchmark_flag: Whether to run benchmark circuits
        data_receiver: GetData instance for execution
    """

    def __init__(self, platform_config: dict, circuit_nums: int, shots: int,
                 benchmark_flag: bool = False):
        """Initialize circuit handler.

        Args:
            platform_config: Platform configuration dictionary
            circuit_nums: Number of circuits to run
            shots: Number of shots per circuit
            benchmark_flag: Whether to run benchmark circuits
        """
        self.platform_config = platform_config
        self.circuit_nums = circuit_nums
        self.shots = shots
        self.benchmark_flag = benchmark_flag
        self.data_receiver = GetData(platform_config['platform'],
                                   platform_config['login_key'],
                                   platform_config['machine_name'],
                                   self.shots,
                                   platform_config['num_qubits'],
                                   platform_config['exp_name'])

    def submit_experiment(self, target_circuit: Union[str, List[str]]) -> None:
        """Submit quantum circuits for execution.

        Args:
            target_circuit: Circuit(s) to execute, either single circuit or list
        """
        if isinstance(target_circuit, str):
            target_circuit = [target_circuit] * self.circuit_nums

        self.data_receiver.submit_experiment(target_circuit, 
                                          benchmark=self.benchmark_flag)

    def receive_results(self) -> None:
        """Collect results from experiment."""
        self.data_receiver.receive_results()

    def save_results(self, save_path: str, simplify: bool = True) -> None:
        """Save experiment results.

        Args:
            save_path: Directory to save results
            simplify: Whether to save only essential columns
        """
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        file_name = f'{self.platform_config["exp_name"]}_{self.shots}shots.csv'
        self.data_receiver.save_results(os.path.join(save_path, file_name), simplify)
